# -*- coding: utf-8 -*-
"""

 焦距可以两种方法：
 1  self.focal_length = rect["Q"][2,3]  # 默认值，一般取立体校正后的重投影矩阵Q中的 Q[2,3]
2  相机6000 is 6000x4000pix and sensor size is 23.5 x 15.6mm hence:
 Fx = fx * W /w = 5192pix * 23.5mm / 6000pix = 20.33mm
  Fy = fy * H /h =  5192pix * 15.6mm / 4000pix = 20.25mm
   正方形像素两者基本相等

class Conf:
    def __init__(self):
        self.cur=os.path.dirname(__file__)
        cali=np.load(os.path.join(self.cur, "cab.npz"))
        # 相机内参
        self.mat=cali['mat']
        # 相机畸变系数:[k1, k2, p1, p2, k3]
        self.dist=cali['dist']
    
    @property
    def focal_length(self):
        Fx=self.mat[0, 0] * self.pix2cm
        return Fx
    
    @property
    def pix2cm(self):
        ccd_w=2.36
        ccd_h=1.58
        img_w=3872
        img_h=2592
        xu=ccd_w / img_w
        yu=ccd_h / img_h
        # 无论x或者y两个方向计算像素长度，基本一样
        return xu
    
    def rectify_distortion(self, p):
        img=cv2.imread(p)
        h, w=img.shape[:2]
        newcameramtx, roi=cv2.getOptimalNewCameraMatrix(self.mat, self.dist, (w, h), 1, (w, h))
        # undistort
        dst=cv2.undistort(img, mtx, dist, None, newcameramtx)
        # crop the image
        x, y, w, h=roi
        dst=dst[y:y + h, x:x + w]
        cv2.imwrite('calibresult.png', dst)


   畸变矫正是上一篇博文的遗留问题，当畸变系数和内外参数矩阵标定完成后，就应该进行畸变的矫正，以达到消除畸变的目的，此其一。

在该系列第一部分的博文中介绍的立体成像原理中提到，要通过两幅图像估计物点的深度信息，就必须在两幅图像中准确的匹配到同一物点，这样才能根据该物点在两幅图像中的位置关系，计算物体深度。为了降低匹配的计算量，两个摄像头的成像平面应处于同一平面。但是，单单依靠严格的摆放摄像头来达到这个目的显然有些困难。立体校正就是利用几何图形变换(Geometric Image Transformation)关系，使得原先不满足上述位置关系的两幅图像满足该条件，此其二。
computes the camera matrices, distortion efficients of left and right cameras.
obtains the transformation from left to right camera, the essential matrix, the fundamental matrix.
take an image as an example to to recitification
"""

import numpy as np
import os, cv2
import glob


# monocular camera calibration

class Cali:
    def __init__(self):
        self.dir_cali="./cali_param"
        os.makedirs(self.dir_cali, exist_ok=True)
        self.param=f"{self.dir_cali}/param.npz"
        self.chess_boards=f"{self.dir_cali}/chess_boards.npz"
        self.rect=f"{self.dir_cali}/rect.npz"
    
    def binocular_calibrate(self):
        cbr=6
        cbc=9
        objpoints=[]
        imgpoints1=[]
        imgpoints2=[]
        criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        objp=np.zeros((cbr * cbc, 3), np.float32)
        objp[:, :2]=np.mgrid[0:cbc, 0:cbr].T.reshape(-1, 2)
        
        # left camera calibration
        images=glob.glob('./cali_img/left/*.jpg')
        gray_shape=None
        for fname in images:
            img=cv2.imread(fname)
            gray=cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            gray_shape=gray.shape[::-1]
            ret, corners=cv2.findChessboardCorners(gray, (cbc, cbr), None)
            if ret==True:
                objpoints.append(objp)
                corners2=cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
                imgpoints1.append(corners2)
        ret, mtx_l, dist_l, rvecs_l, tvecs_l=cv2.calibrateCamera(objpoints, imgpoints1, gray_shape, None, None)
        
        # right camera calibration
        images=glob.glob('./cali_img/right/*.jpg')
        for fname in images:
            img=cv2.imread(fname)
            gray=cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            gray_shape=gray.shape[::-1]
            ret, corners=cv2.findChessboardCorners(gray, (cbc, cbr), None)
            if ret==True:
                corners2=cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
                imgpoints2.append(corners2)
        ret, mtx_r, dist_r, rvecs_r, tvecs_r=cv2.calibrateCamera(objpoints, imgpoints2, gray_shape, None, None)
        
        # binocular camera calibration
        ret, mtx_l, dist_l, mtx_r, dist_r, R, T, E, F=cv2.stereoCalibrate(objpoints, imgpoints1, imgpoints2, mtx_l,
                                                                          dist_l, mtx_r, dist_r, gray_shape)
        
        np.savez(self.param, ret=ret, mtx_l=mtx_l, mtx_r=mtx_r, dist_l=dist_l, dist_r=dist_r, R=R, T=T,
                 gray_shape=gray_shape)
        np.savez(self.chess_boards, objpoints=objpoints, imgpoints1=imgpoints1, imgpoints2=imgpoints2)
        
        print('intrinsic matrix of left camera=\n', mtx_l)
        print('intrinsic matrix of right camera=\n', mtx_r)
        print('distortion coefficients of left camera=\n', dist_l)
        print('distortion coefficients of right camera=\n', dist_r)
        print('Transformation from left camera to right:\n')
        print('R=\n', R)
        print('\n')
        print('T=\n', T)
        print('\n')
        print('Reprojection Error=\n', ret)
        print('gray shape=\n', gray_shape)
    
    # stereo rectification
    def binocular_rect(self):
        param=np.load(self.param)
        # 相机内参
        mtx_l=param['mtx_l']
        dist_l=param['dist_l']
        mtx_r=param['mtx_r']
        dist_r=param['dist_r']
        R=param['R']
        T=param['T']
        gray_shape=param['gray_shape']
        R1, R2, P1, P2, Q, ROI1, ROI2=cv2.stereoRectify(mtx_l, dist_l, mtx_r, dist_r, gray_shape, R, T, flags=0,
                                                        alpha=-1)
        
        # undistort rectifying mapping
        map1_l, map2_l=cv2.initUndistortRectifyMap(mtx_l, dist_l, R1, P1, gray_shape, cv2.CV_16SC2)
        map1_r, map2_r=cv2.initUndistortRectifyMap(mtx_r, dist_r, R2, P2, gray_shape, cv2.CV_16SC2)
        np.savez(self.rect, R1=R1, R2=R2, P1=P1, P2=P2, Q=Q, ROI1=ROI1, ROI2=ROI2, map1_l=map1_l, map1_r=map1_r,
                 map2_l=map2_l, map2_r=map2_r)
    
    def rect_pair(self):
        os.makedirs("./rectifyresult/left", exist_ok=True)
        os.makedirs("./rectifyresult/right", exist_ok=True)
        rc=np.load(self.rect)
        map1_l=rc['map1_l']
        map1_r=rc['map1_r']
        map2_r=rc['map2_r']
        map2_l=rc['map2_l']
        # undistort the original image, take img#3 as an example
        pics=os.listdir("./img/left")
        for pic in pics:
            
            left3=cv2.imread(f"./img/left/{pic}")
            dst_l=cv2.remap(left3, map1_l, map2_l, cv2.INTER_LINEAR)
            sucl=cv2.imwrite(f'./rectifyresult/left/{pic}(', dst_l)
            
            right3=cv2.imread(f"./img/right/{pic}")
            dst_r=cv2.remap(right3, map1_r, map2_r, cv2.INTER_LINEAR)
            sucr=cv2.imwrite(f'./rectifyresult/right/{pic}', dst_r)
            if sucl and sucr:
                print("sucess rect pic", pic)


if __name__=='__main__':
    c=Cali()
    c.binocular_calibrate()
    c.binocular_rect()
    c.rect_pair()
